import numpy as np
import pandas as pd

# 读取数据
df = pd.read_excel('stock.xlsx', dtype={'code': 'str'})

# 将 timeToMarket 列转换为字符串，然后提取年份
df['year'] = df['timeToMarket'].astype(str).str[:4]

# 计算年份对应的股票数量
yearnum = df.groupby('year').size()

# 计算PE均值
pe_mean = df['pe'].mean()

# 计算tvalue
df['tvalue'] = 4 * df['esp'] - df['pe'] * df['totals']
weighted_pe = np.sum(df['pe'] * df['tvalue']) / df['tvalue'].sum()

# 将索引转换为字符串列
df['board'] = df.index.astype(str).str[:2]

# 计算板块的PE均值和股票数
board_stats = df.groupby('board')['pe'].agg([('pe均值', 'mean'), ('股票数', 'count')])

# 打印结果
print("年份对应的股票数量：\n", yearnum)
print("PE均值：", pe_mean)
print("加权PE均值：", weighted_pe)
print("板块统计：\n", board_stats)
